package com.bins.cloud.gateway.limiter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;

import java.util.Arrays;
import java.util.List;

@Slf4j
@Component
public class CustomRedisRateLimiter {

    private final RedisTemplate<String, Long> redisTemplate;

    private final RedisScript<List<Long>> redisScript;

    /**
     * @param redisScript 该RedisScript对象将会自动注入进来，该对象使用的正是上面介绍的request-rate-limiter.lua脚本
     */
    public CustomRedisRateLimiter(RedisTemplate redisTemplate, RedisScript<List<Long>> redisScript) {
        this.redisTemplate = redisTemplate;
        this.redisScript = redisScript;
    }

    public boolean isAllowed(String key, int replenishRate, int burstCapacity) {
        try {
            String luaScriptPath = "token-bucket.lua";
            Resource scriptSource = new ClassPathResource(luaScriptPath);
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
            redisScript.setScriptSource(new ResourceScriptSource(scriptSource));
            redisScript.setResultType(Long.class);
            List<String> keys = Arrays.asList(key + "tokens", key + "timestamp");
            long now = System.currentTimeMillis();
            this.redisTemplate.setKeySerializer(new StringRedisSerializer());
            this.redisTemplate.setValueSerializer(new StringRedisSerializer());
            Long result = this.redisTemplate.execute(redisScript, keys, replenishRate + "", burstCapacity + "", now + "", 1 + "");
            log.info("result:{}", result);
            if (ObjectUtils.isEmpty(result) || result == 0) {
                return false;
            } else {
                return true;
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            return true;
        }

    }
}